#!/usr/bin/env python
# coding: utf-8

# In[4]:


import numpy as np
import pandas as pd
from pytorch_tabnet.tab_model import TabNetClassifier
from sklearn.metrics import roc_auc_score
from sklearn.preprocessing import LabelEncoder

np.random.seed(0)

import scipy

from matplotlib import pyplot as plt

plt.show()

# In[5]:


# 多GPU环境下，指定使用第2张GPU
import os

os.environ['CUDA_VISIBLE_DEVICES'] = f"0"

# In[6]:


import torch

torch.__version__

# # Download census-income dataset

# In[7]:


# 获取《抑郁症》数据集，out=数据集路径
dataset_path = ["./features/sleepEDF-cassette_features.xlsx", "./features/hospital_features.xlsx"]

# In[9]:

# 读取所有 Excel 文件并存储在一个列表中
dfs = [pd.read_excel(file) for file in dataset_path]
# 垂直拼接
train = pd.concat(dfs, axis=0, ignore_index=True)
target = 'isMDD'  # 设置目标列名称（分类任务的目标变量）
# 如果还没划分训练集、验证集、测试集，就新增Set列（用于标记每一行属于哪个集合）
# 概率80%，10%，10%
if "Set" not in train.columns:
    train["Set"] = np.random.choice(["train", "valid", "test"], p=[.8, .1, .1], size=(train.shape[0],))

train_indices = train[train.Set == "train"].index  # 存训练集的行的索引
valid_indices = train[train.Set == "valid"].index
test_indices = train[train.Set == "test"].index

# # Simple preprocessing
#
# Label encode categorical features and fill empty cells.
#
# 对分类特征进行标签编码，并填充空单元格。

# In[10]:


nunique = train.nunique()  # 获取每列的不同值的数量
types = train.dtypes  # 获取每列的数据类型

categorical_columns = []  # 存储分类特征列的列名（原本就是分类特征列，或因为不同值太少当成分类特征列）
categorical_dims = {}  # 存储每个分类特征列中的类别数
for col in train.columns:
    if types[col] == 'object' or nunique[col] < 200:  # 如果列的数据类型是object（分类特征）或不同值数量少于200
        print(col, train[col].nunique())  # 打印列名和不同值数量
        train[col] = train[col].fillna(-1)  # 缺失值填充为字符串 "VV_likely"
        # 创建一个“标签编码器”，将该列的分类特征（字符串或分类值）转换为数值
        l_enc = LabelEncoder()
        train[col] = l_enc.fit_transform(train[col].values)

        categorical_columns.append(col)
        categorical_dims[col] = len(l_enc.classes_)
    else:  # 如果列的数据类型是数值，缺失值填充为训练集部分的均值
        train.fillna(train.loc[train_indices, col].mean(), inplace=True)

    # In[11]:

# 保持目标列为原始数值类型
train[target] = train[target].astype(int)

# # Define categorical features for categorical embeddings
# 定义分类特征的列索引、类别数，为后续将特征类别从独热编码转变成嵌入编码做准备

# In[12]:


unused_feat = ['Set']  # 将 Set 列标记为不使用的特征列（因为是分训练、验证、测试集用的）

# 构建一个包含所有用于模型训练的特征列的列表 features（排除了 Set列、target列）
features = [col for col in train.columns if col not in unused_feat + [target]]

# cat_idxs列表 存储所有使用的特征列中，是分类特征的列索引
cat_idxs = [i for i, f in enumerate(features) if f in categorical_columns]

# cat_dims列表 存储所有使用的特征列中，是分类特征的列的类别数
cat_dims = [categorical_dims[f] for i, f in enumerate(features) if f in categorical_columns]
# cat_dims相比于categorical_dims，减少了 不使用列，且不是字典是数组


# # Grouped features 特征组
#
# You can now specify groups of feature which will share a common attention.
#
# 你现在可以指定一组特征，使它们共享一个共同的注意力机制。
#
# This may be very usefull for features comming from a same preprocessing technique like PCA for example.
#
# 这对于来自相同预处理技术（例如 PCA）生成的特征可能非常有用。

# In[13]:


len(features)

# In[14]:


# 对相关的特征分进一组，同组特征共享一个注意力机制
# 更好地捕获分组内特征之间的关系
# 模型能够聚焦于某些更重要的特征组合
grouped_features = [[0, 1], [2, 3, 4], [5, 6]]

# # Network parameters
# TabNet 的初始化配置（超参数的设置、模型的定义）

# In[15]:


tabnet_params = {"cat_idxs": cat_idxs,
                 "cat_dims": cat_dims,
                 "cat_emb_dim": 2,  # 每个分类特征类别的嵌入向量维度（每个分类特征将映射到一个 2 维的嵌入空间）
                 "optimizer_fn": torch.optim.Adam,  # 指定使用torch.optim.Adam 优化器
                 "optimizer_params": dict(lr=2e-2),  # 设置优化器超参数，学习率lr=2e-2(即 0.02)，表示模型的学习步长

                 # 设置学习率调度器的参数（每 50 个训练周期调整一次学习率、每次调整变成原来的 90%）
                 "scheduler_params": {"step_size": 50,  # how to use learning rate scheduler
                                      "gamma": 0.9},
                 # 设置学习率调度的策略（根据步长调整学习率）
                 "scheduler_fn": torch.optim.lr_scheduler.StepLR,
                 # 设置用于注意力机制的掩码类型（选平滑稀疏化、而非完全稀疏化）
                 "mask_type": 'entmax',  # "sparsemax"
                 # 设置特征分组
                 "grouped_features": grouped_features
                 }

# 用前面设置好的超参数，配置 TabNet 的分类器
clf = TabNetClassifier(**tabnet_params
                       )

# # Training

# In[16]:


# 提取特征矩阵、目标值向量
X_train = train[features].values[train_indices]
y_train = train[target].values[train_indices]

X_valid = train[features].values[valid_indices]
y_valid = train[target].values[valid_indices]

X_test = train[features].values[test_indices]
y_test = train[target].values[test_indices]

# In[17]:


# 设置模型的最大迭代次数
# 在 CI 环境中，检查代码逻辑是否正确，而不需要完整训练
# 在本地环境中，才多次迭代，完整训练
max_epochs = 50 if not os.getenv("CI", False) else 2

# In[18]:


# 引入 TabNet 的增强策略，ClassificationSMOTE 类用于生成额外的少数类样本，提升模型对少数类的预测性能
# 使增强样本占整个训练批次的 20%
from pytorch_tabnet.augmentations import ClassificationSMOTE

aug = ClassificationSMOTE(p=0.2)

# In[19]:


# This illustrates the behaviour of the model's fit method using Compressed Sparse Row matrices
# 这段代码展示了如何使用 Compressed Sparse Row (CSR) 矩阵 格式的输入数据来训练 TabNetClassifier 模型

# 将训练、验证特征矩阵转换为 CSR 矩阵 格式
sparse_X_train = scipy.sparse.csr_matrix(X_train)  # Create a CSR matrix from X_train
sparse_X_valid = scipy.sparse.csr_matrix(X_valid)  # Create a CSR matrix from X_valid

# Fitting the model
# 这段代码展示了模型 fit 方法的一次基本使用
clf.fit(
    X_train=sparse_X_train, y_train=y_train,

    # 设置验证（验证数据、验证集名称、评估的指标）
    eval_set=[(sparse_X_train, y_train), (sparse_X_valid, y_valid)],
    eval_name=['train', 'valid'],
    eval_metric=['auc'],

    # 若验证集性能在 patience 个轮次内无提升，则停止训练
    max_epochs=max_epochs, patience=20,

    # 设置批量大小（每次从训练集中取 1024 条样本进行一次模型迭代）
    # 虚拟批量大小【当 GPU 显存不足以直接处理 batch_size 大小的样本时，模型分多次加载更小的批次（virtual_batch_size）。不影响效果，只影响时间】
    batch_size=171, virtual_batch_size=50,

    num_workers=0,  # 数据加载器的线程数（0 表示单线程）
    weights=1,  # 样本权重（1 表示均等权重）
    drop_last=False,
    # 设置使用的数据增强方法
    augmentations=aug,  # aug, None
)

# In[20]:


# This illustrates the warm_start=False behaviour
# 展示非热启动下，是否计算特征重要性不会影响模型性能
# 非热启动（每次启动模型训练，都是从头开始）
# 热启动（基于之前已经训练过的参数值）
save_history = []

# Fitting the model without starting from a warm start nor computing the feature importance
# 在既不热启动，也不计算特征重要性的情况下，训练模型2次
for _ in range(2):
    clf.fit(
        X_train=X_train, y_train=y_train,
        eval_set=[(X_train, y_train), (X_valid, y_valid)],
        eval_name=['train', 'valid'],
        eval_metric=['auc'],
        max_epochs=max_epochs, patience=20,
        batch_size=1024, virtual_batch_size=128,
        num_workers=0,
        weights=1,
        drop_last=False,
        augmentations=aug,  # aug, None
        # 不计算特征重要性
        compute_importance=False
    )
    # 记录2次训练出模型的"valid_auc"参数
    save_history.append(clf.history["valid_auc"])

# 验证2次"valid_auc"参数相同
assert (np.all(np.array(save_history[0] == np.array(save_history[1]))))

save_history = []  # Resetting the list to show that it also works when computing feature importance

# Fitting the model without starting from a warm start but with the computing of the feature importance activated
# 在不热启动，但计算特征重要性的情况下，训练模型2次
for _ in range(2):
    clf.fit(
        X_train=X_train, y_train=y_train,
        eval_set=[(X_train, y_train), (X_valid, y_valid)],
        eval_name=['train', 'valid'],
        eval_metric=['auc'],
        max_epochs=max_epochs, patience=20,
        batch_size=1024, virtual_batch_size=128,
        num_workers=0,
        weights=1,
        drop_last=False,
        augmentations=aug,  # aug, None
        # 计算特征重要性
        compute_importance=True  # True by default so not needed
    )
    save_history.append(clf.history["valid_auc"])

assert (np.all(np.array(save_history[0] == np.array(save_history[1]))))

# In[21]:


# plot losses
plt.plot(clf.history['loss'])

# In[22]:


# plot auc
plt.plot(clf.history['train_auc'])
plt.plot(clf.history['valid_auc'])

# In[23]:


# plot learning rates
plt.plot(clf.history['lr'])

# ## Predictions

# In[36]:


# 模型clf 对测试集预测
preds = clf.predict_proba(X_test)
# 计算测试集的 AUC（反映泛化性能）
test_auc = roc_auc_score(y_score=preds[:, 1], y_true=y_test)

preds_valid = clf.predict_proba(X_valid)
# 计算验证集的 AUC（反映是否过拟合）
valid_auc = roc_auc_score(y_score=preds_valid[:, 1], y_true=y_valid)

print(f"BEST VALID SCORE FOR {dataset_path} : {clf.best_cost}")  # 验证集AUC
print(f"FINAL TEST SCORE FOR {dataset_path} : {test_auc}")  # 测试集AUC

# In[25]:


# check that best weights are used
# 检查，当前模型所使用的权重 = 对应验证集 AUC 最优时的权重
assert np.isclose(valid_auc, np.max(clf.history['valid_auc']), atol=1e-6)

# In[26]:


# 展示clf 对测试集预测结果的数组
clf.predict(X_test)

# # Save and load Model
# 保存并加载模型
#

# In[27]:


# save tabnet model
saving_path_name = "./tabnet_model_test_MDD_1"
saved_filepath = clf.save_model(saving_path_name)

# In[28]:


# define new model with basic parameters and load state dict weights
loaded_clf = TabNetClassifier()  # 初始化一个新的TabNet模型
loaded_clf.load_model(saved_filepath)  # 加载保存的权值、配置

# In[29]:


# 用加载的模型预测，检查AUC
loaded_preds = loaded_clf.predict_proba(X_test)
loaded_test_auc = roc_auc_score(y_score=loaded_preds[:, 1], y_true=y_test)

print(f"FINAL TEST SCORE FOR {dataset_path} : {loaded_test_auc}")

# In[30]:


assert (test_auc == loaded_test_auc)

# In[31]:


loaded_clf.predict(X_test)

# # Global explainability : feat importance summing to 1
# 全局可解释性

# In[32]:


# 全局特征重要性掩码 = Mglobal数组（所有样本的聚合特征重要性掩码 Magg[ i ] 的加权平均）
clf.feature_importances_

# # Local explainability and masks
# 局部可解释性与掩码

# In[50]:


# explain_matrix(3267, 14) = 【3个步骤加权平均】每个样本对14个特征的选择权重（聚合特征重要性掩码 Magg[ i ]）
# masks(3，3267, 14) = 【3个决策步骤中】每个样本对14个特征的选择权重（特征重要性掩码 M[ i ]）
explain_matrix, masks = clf.explain(X_test)
print(explain_matrix.shape)
print(len(masks))
print(masks[1].shape)

# In[34]:


# 3个子图，对应3个决策步骤
fig, axs = plt.subplots(1, 3, figsize=(20, 20))

# 显示前50个样本，3个决策步骤，每个样本对14个特征的选择权重（颜色越亮权值越大）
for i in range(3):
    axs[i].imshow(masks[i][:50])
    axs[i].set_title(f"mask {i}")
    axs[i].set_xticklabels(labels=features, rotation=45)

# # XGB
